import onnxruntime as ort
import numpy as np
import cv2


def resize_and_pad_image(image, target_size):
    # 获取图像的高度和宽度
    height, width = image.shape[:2]

    # 计算缩放比例
    scale = target_size / max(height, width)
    new_height, new_width = int(height * scale), int(width * scale)

    # 等比例缩放图像
    resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)

    # 计算填充量
    pad_vert = (target_size - new_height) // 2
    pad_horz = (target_size - new_width) // 2

    # 应用填充
    padded_image = cv2.copyMakeBorder(resized_image, pad_vert, pad_vert, pad_horz, pad_horz,
                                      cv2.BORDER_CONSTANT, value=[0, 0, 0])  # 使用黑色进行填充

    return padded_image


class ONNXClassifier:
    def __init__(self, model_path):
        self.session = ort.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

    def preprocess_image(self, image, target_size=224):
        # 调用 resize_and_pad_image 方法
        processed_image = resize_and_pad_image(image, target_size)
        processed_image = cv2.resize(processed_image, (target_size, target_size))
        # 转换图像为模型需要的格式（例如：NCHW）
        processed_image = np.transpose(processed_image, (2, 0, 1)).astype(np.float32) / 255.0
        # 添加批次维度
        processed_image = np.expand_dims(processed_image, axis=0)
        return processed_image

    def predict(self, image):
        # 预处理图像
        try:
            preprocessed_image = self.preprocess_image(image)
            # 进行推理
            outputs = self.session.run([self.output_name], {self.input_name: preprocessed_image})
            # 获取预测结果（概率分布）
            probabilities = outputs[0][0]  # 假设输出的shape是[1, num_classes]

        except Exception:
            return 0

        # 检查第一个类别的概率是否小于0.7
        if probabilities[0] < 0.85:
            # 忽略第一个类别
            probabilities[0] = 0
            return np.argmax(probabilities)  # 返回除第一个类别外概率最高的类别索引
        else:
            return np.argmax(probabilities)  # 返回概率最高的类别索引


# 使用示例
if __name__ == "__main__":
    classifier = ONNXClassifier("model.onnx")  # 替换为你的模型路径
    image_path = "/Users/tunm/datasets/person_uniform/classify/2/db618461-a341-4798-a46c-465a96976ff6.jpg"  # 替换为你的图像路径
    image = cv2.imread(image_path)
    prediction = classifier.predict(image)
    print("Predicted class:", prediction)
